#include "target_tracker.h"
#include <algorithm>

TargetTracker::TargetTracker() : max_lost_count_(5), bullet_speed_(25.0) {} // 25 cm/s

void TargetTracker::initKalmanFilter(cv::KalmanFilter& kf) {
    kf.init(4, 2, 0); // 4状态量(x,y,vx,vy)，2观测量(x,y)
    
    // 转移矩阵 [1 0 dt 0; 0 1 0 dt; 0 0 1 0; 0 0 0 1]
    kf.transitionMatrix = (cv::Mat_<float>(4, 4) << 
        1, 0, 1, 0,
        0, 1, 0, 1,
        0, 0, 1, 0,
        0, 0, 0, 1);
    
    // 测量矩阵 [1 0 0 0; 0 1 0 0]
    cv::setIdentity(kf.measurementMatrix);
    
    // 过程噪声协方差
    cv::setIdentity(kf.processNoiseCov, cv::Scalar::all(1e-2));
    
    // 测量噪声协方差
    cv::setIdentity(kf.measurementNoiseCov, cv::Scalar::all(1e-1));
    
    // 后验误差协方差
    cv::setIdentity(kf.errorCovPost, cv::Scalar::all(1));
}

cv::Point2f TargetTracker::applyKalmanFilter(cv::KalmanFilter& kf, const cv::Point2f& measurement) {
    // 预测
    cv::Mat prediction = kf.predict();
    
    // 更新
    cv::Mat measurement_mat(2, 1, CV_32F);
    measurement_mat.at<float>(0) = measurement.x;
    measurement_mat.at<float>(1) = measurement.y;
    kf.correct(measurement_mat);
    
    return cv::Point2f(prediction.at<float>(0), prediction.at<float>(1));
}

void TargetTracker::update(const std::vector<Armor>& detected_armors) {
    // 简化版跟踪逻辑
    // 实际应该使用更复杂的数据关联算法
    
    // 更新现有目标
    for (auto& target : targets_) {
        target.lost_count++;
    }
    
    // 关联新检测
    for (const auto& armor : detected_armors) {
        bool matched = false;
        for (auto& target : targets_) {
            double distance = cv::norm(armor.getCenter() - target.armor.getCenter());
            if (distance < 50) { // 距离阈值
                target.armor = armor;
                target.lost_count = 0;
                target.tracked_count++;
                target.trajectory.push_back(armor.getCenter());
                
                // 应用卡尔曼滤波
                applyKalmanFilter(target.kf, armor.getCenter());
                
                matched = true;
                break;
            }
        }
        
        if (!matched) {
            TrackedTarget new_target;
            new_target.armor = armor;
            new_target.lost_count = 0;
            new_target.tracked_count = 1;
            new_target.trajectory.push_back(armor.getCenter());
            
            initKalmanFilter(new_target.kf);
            new_target.kf.statePost.at<float>(0) = armor.getCenter().x;
            new_target.kf.statePost.at<float>(1) = armor.getCenter().y;
            new_target.kf.statePost.at<float>(2) = 0;
            new_target.kf.statePost.at<float>(3) = 0;
            
            targets_.push_back(new_target);
        }
    }
    
    // 移除丢失的目标
    targets_.erase(
        std::remove_if(targets_.begin(), targets_.end(),
            [this](const TrackedTarget& target) {
                return target.lost_count > max_lost_count_;
            }),
        targets_.end()
    );
}

Armor TargetTracker::selectTarget() {
    if (targets_.empty()) {
        return Armor(); // 返回默认构造的 Armor
    }
    // Scoring strategy: prefer long-tracked, recently-seen and near-center targets
    double best_score = -1e9;
    int best_idx = 0;
    // Use approximate frame center; ideally should be configurable
    cv::Point2f frame_center(640.0f/2.0f, 480.0f/2.0f);
    for (size_t i = 0; i < targets_.size(); ++i) {
        const auto &t = targets_[i];
        double score = 0.0;
        score += t.tracked_count * 2.0; // stability
        score -= t.lost_count * 1.5; // penalize lost
        double dist_center = cv::norm(t.trajectory.back() - frame_center);
        score -= dist_center * 0.01; // prefer near-center
        // small bonus for more recent tracking
        score += std::min(10, t.tracked_count);

        if (score > best_score) {
            best_score = score;
            best_idx = static_cast<int>(i);
        }
    }

    return targets_[best_idx].armor;
}

cv::Point2f TargetTracker::predictHitPoint(const Armor& target, double bullet_speed) {
    if (targets_.empty()) {
        return target.getCenter();
    }
    
    // 找到对应的跟踪目标
    for (auto& tracked_target : targets_) {
        if (cv::norm(tracked_target.armor.getCenter() - target.getCenter()) < 10) {
            // Predict using current Kalman state (px,py,vx,vy)
            cv::Mat state = tracked_target.kf.statePost;
            float px = state.at<float>(0);
            float py = state.at<float>(1);
            float vx = state.at<float>(2);
            float vy = state.at<float>(3);

            // Time-of-flight (s): use armor distance if available
            double tof = 0.2; // default fallback
            double armor_dist_mm = tracked_target.armor.getDistance();
            if (armor_dist_mm > 0.0) {
                double dist_m = armor_dist_mm / 1000.0; // mm -> m
                double bullet_m_s = (bullet_speed) / 100.0; // cm/s -> m/s
                if (bullet_m_s > 1e-3) tof = dist_m / bullet_m_s;
            }

            cv::Point2f predict_pos(px + vx * static_cast<float>(tof), py + vy * static_cast<float>(tof));
            return predict_pos;
        }
    }
    
    return target.getCenter();
}

void TargetTracker::drawPrediction(cv::Mat& frame, const cv::Point2f& predict_point) {
    cv::circle(frame, predict_point, 8, cv::Scalar(0, 0, 255), 3); // 红色圆圈
    cv::circle(frame, predict_point, 3, cv::Scalar(0, 0, 255), -1);
}